import os
import json
import cv2
from lmk68 import trans_dlib_68
from face_similarity import procrustes_similarity
import numpy as np
def getFaceDataset():
    baselist=[]
    path = "F:\data\ThreeBallFace"
    fileList = os.listdir(path)
    for file in fileList:
        filepath = os.path.join(path, file)
        if os.path.isdir(filepath):
            baselist.append(filepath)
    return baselist
def alignment(baseData,faceData):
    # 假设src_pts和dst_pts是Nx2的numpy数组
    src_pts = np.array(faceData, dtype=np.float32)
    dst_pts = np.array(baseData, dtype=np.float32)
    # 计算仿射变换矩阵
    M, _ = cv2.estimateAffine2D(src_pts, dst_pts)
    # 应用变换到源关键点
    aligned_pts = cv2.transform(src_pts.reshape(1, -1, 2), M).reshape(-1, 2)
    return aligned_pts,M
face="F:\data\ThreeBallFace/223_393-1655375351937_whiteFace_face.jpg"
face_json=face.replace(".jpg",".json")
baselist=getFaceDataset()
maxScore=0
maxpath=None
for baseFace in baselist:
    blist=os.listdir(baseFace)
    bimg,bjson=None,None
    for dir in ["whiteFace_face.jpg","whiteLeft_left.jpg","whiteRight_right.jpg"]:
        if dir in face:
            bimg = [img for img in blist if "whiteFace_face.jpg" in img][0]
            bjson=bimg.replace(".jpg",".json")
            break
    bimgpath=os.path.join(baseFace,bimg)
    bjsonpath=os.path.join(baseFace,bjson)
    basedata,facedata=None,None
    with open(bjsonpath, 'r') as file:
        basedata = json.load(file)
    with open(face_json, 'r') as file:
        facedata = json.load(file)
    alignmentface,alignmentM=alignment(basedata["landmarks"],facedata["landmarks"])
    alignment_data = alignmentface.tolist()
    similarity = procrustes_similarity(basedata["landmarks"], alignment_data)
    if similarity>maxScore:
        maxScore=similarity
        maxpath=baseFace
        print(similarity)
print(maxpath)
    # faceimg=cv2.imread(face)
    # h1,w1,c=faceimg.shape
    # baseimg=cv2.imread(bimgpath)
    # h2,w2,c=baseimg.shape
    # print(len(basedata["landmarks"]))
    #
    # img68 = baseimg.copy()
    # for idx in range(106):
    #     base_pt=basedata["landmarks"][idx]
    #     face_pt=facedata["landmarks"][idx]
    #     ali_pt=alignment_data[idx]
    #     cv2.circle(baseimg, (int(base_pt[0]),int(base_pt[1])), 3, (0, 255, 0), -1)
    #     cv2.circle(faceimg, (int(face_pt[0]), int(face_pt[1])), 3, (0, 255, 0), -1)
    #
    #     cv2.circle(baseimg, (int(ali_pt[0]), int(ali_pt[1])), 3, (0, 0, 255), -1)
    # pts68=trans_dlib_68(alignment_data)
    #
    # for idx in range(68):
    #     base_pt=pts68[idx]
    #     cv2.circle(img68, (int(base_pt[0]),int(base_pt[1])), 3, (0, 255, 0), -1)
    # showface = cv2.resize(faceimg, (w1//2, h1//2), interpolation=cv2.INTER_AREA)
    # showbase = cv2.resize(baseimg, (w2 // 2, h2 // 2), interpolation=cv2.INTER_AREA)
    # showbase68 = cv2.resize(img68, (w2 // 2, h2 // 2), interpolation=cv2.INTER_AREA)
    # cv2.imshow("faceimg",showface)
    # cv2.imshow("baseimg",showbase)
    # cv2.imshow("showbase68", showbase68)
    # cv2.waitKey()


